//
// Created by wya on 2021/5/10.
//
#include "polarization.h"
#include <iostream>
#include <opencv2/core.hpp>
//#include <eigen3/Eigen/Dense>


#include <math.h>
#include "imageProcess.h"


using namespace arma;

using namespace NfP;
void NfP::Polarized::NormalOfPolarized(const std::string& path, const std::string& save_path,NfP::nfp_parameters* p)
{

    clock_t start = clock();
    NfP::Polarized::Initial(path,p);
    NfP::Polarized::Compute_polarization(p);


    NfP::Polarized::rho_diffuse(p);
    NfP::Polarized::rho_spec(p);
    /*******show phi map****/
    cv::Mat phi;
    cv::transpose(cv::Mat_<double>(static_cast<int>(p->phi.n_cols),
                                   static_cast<int>(p->phi.n_rows),
                                   const_cast<double*>(p->phi.memptr())),
                  phi);
    cv::Mat phi_gary,phi_color;
    phi.convertTo(phi_gary,CV_8UC1,255);
    cv::applyColorMap(phi_gary, phi_color, cv::COLORMAP_HOT);
    cv::imshow("phi",phi_gary);

    /*******show normal map****/
    NfP::Polarized::polar2normal(p->theta,p->phi,&p->normal);

    cv::Mat normal_x=cv::Mat::ones(p->rows,p->cols,CV_64FC1);
    cv::Mat normal_y=cv::Mat::ones(p->rows,p->cols,CV_64FC1);
    cv::Mat normal_z=cv::Mat::ones(p->rows,p->cols,CV_64FC1);
    cv::Mat normal=cv::Mat::ones(p->rows,p->cols,CV_64FC3);


    cv::transpose(cv::Mat_<double>(static_cast<int>(p->normal.slice(0).n_cols),
                                   static_cast<int>(p->normal.slice(0).n_rows),
                                   const_cast<double*>(p->normal.slice(0).memptr())),
                  normal_x);
    cv::transpose(cv::Mat_<double>(static_cast<int>(p->normal.slice(1).n_cols),
                                   static_cast<int>(p->normal.slice(1).n_rows),
                                   const_cast<double*>(p->normal.slice(1).memptr())),
                  normal_y);
    cv::transpose(cv::Mat_<double>(static_cast<int>(p->normal.slice(2).n_cols),
                                   static_cast<int>(p->normal.slice(2).n_rows),
                                   const_cast<double*>(p->normal.slice(2).memptr())),
                  normal_z);



    cv::Mat in[]={normal_x,normal_y,normal_z};
    int fromTo1[] = {0,0, 1,1, 2,2};
    cv::mixChannels(in,3,&normal,1,fromTo1,3);

    cv::Mat normal_BGR=cv::Mat::ones(p->rows,p->cols,CV_8UC3);
    normal.convertTo(normal_BGR,CV_8UC3,255,0);

    cv::Mat normal_RGB=cv::Mat::ones(p->rows,p->cols,CV_8UC3);
    cv::cvtColor(normal_BGR,normal_RGB,cv::COLOR_BGR2RGB);

    cv::Mat mask = cv::imread(path+"/mask.png",cv::IMREAD_UNCHANGED);
    cv::Mat img2=cv::Mat::ones(p->rows,p->cols,CV_8UC3);
    normal_RGB.copyTo(img2, mask);
    clock_t ends = clock();
    cout <<"Running Time (s): "<<(double)(ends - start)/ CLOCKS_PER_SEC << endl;
    cv::imshow("normal", img2);
    cv::waitKey(0);

    cv::imwrite(save_path+"/normal.png",img2);

}



void NfP::Polarized::Initial(const std::string& path, NfP::nfp_parameters* p)
{

    cv::Mat image0 = cv::imread(path+"/000000I0.png",cv::IMREAD_GRAYSCALE);
    cv::Mat image1 = cv::imread(path+"/000000I45.png",cv::IMREAD_GRAYSCALE);
    cv::Mat image2 = cv::imread(path+"/000000I90.png",cv::IMREAD_GRAYSCALE);
    cv::Mat mask =  cv::imread(path+"/mask.png",cv::IMREAD_GRAYSCALE);
    cv::Mat spec = cv::imread(path+"/spec.png",cv::IMREAD_GRAYSCALE);
    cv::Mat img,images;
    cv::hconcat(image0,image1,img);
    cv::hconcat(img,image2,images);
    cv::namedWindow("images", cv::WINDOW_AUTOSIZE);
    cv::imshow("images gray",images);
    cv::Mat_<double> images_normalized;
    images.convertTo(images_normalized,CV_64FC1,1.0/255,0);

    p->rows=mask.rows;
    p->cols=mask.cols;

    p->Images.set_size(images_normalized.rows,images_normalized.cols);
    p->Iun.set_size(mask.rows,mask.cols);
    p->phi.set_size(mask.rows,mask.cols);
    p->theta.set_size(mask.rows,mask.cols);
    p->rho.set_size(mask.rows,mask.cols);

    p->mask.set_size(mask.rows,mask.cols);
    p->spec.set_size(mask.rows,mask.cols);
    p->angles(0)=0.0;
    p->angles(1)=0.785;
    p->angles(2)=1.570;
    p->n=1.5;

    p->normal.set_size(mask.rows,mask.cols,3);

    cvMat2armaMat_d(images_normalized,p->Images.memptr());  //Mat_ is different from Mat in read element
    cvMat2armaMat(mask,p->mask.memptr());
    cvMat2armaMat(spec,p->spec.memptr());
//    p->mask.print();

    int32_t trueCount=0;
    int32_t trueCount_=0;
    //A mask has been provided so select only foreground pixels
    for (int j = 0; j < p->cols; j++) {
        for (int k = 0; k < p->rows; k++) {
            if (p->mask(k,j)>0){
                trueCount++;
            }
            if (p->spec(k,j)>0){
                trueCount_++;
            }
        }
    }
    p->mask_rows.set_size(trueCount);
     p->mask_cols.set_size(trueCount);
    p->spec_rows.set_size(trueCount_);
    p->spec_cols.set_size(trueCount_);
    trueCount=0;
    trueCount_=0;
    for (int j = 0; j < p->cols; j++) {
        for (int k = 0; k < p->rows; k++) {
            if (p->mask(k,j)){
                p->mask_cols(trueCount)=j;
                p->mask_rows(trueCount)=k;
                trueCount++;
            }
            if (p->spec(k,j)){
                p->spec_cols(trueCount_)=j;
                p->spec_rows(trueCount_)=k;
                trueCount_++;
            }
        }
    }
}




/* Polarization information is calculated on a per-pixel basis
 *Inputs:
    images - rows by cols by nimages matrix of input images
    angles - vector of polariser angles, length nimages
    mask   - binary foreground mask (true for foreground)
  Outputs:
    rho    - rows by cols matrix containing degree of polarisation
    phi    - rows by cols matrix containing phase angle
    Iun    - rows by cols matrix containing unpolarised intensity
 * */

void NfP::Polarized::Compute_polarization(NfP::nfp_parameters* p )
{


    mat I(p->mask_cols.n_elem,p->image_nums,fill::zeros);
    for (int i = 0; i < p->image_nums; ++i) {
        for (int32_t j = 0; j < p->mask_cols.n_elem ; ++j) {
            I(j,i) = p->Images(p->mask_rows(j),p->mask_cols(j)+i*p->cols) ;
        }
    }



    /* Fast linear method - this seems to produce almost identical results
     to the nonlinear method but is much faster*/
    vec an2 =2*p->angles;
    vec cos_= cos(an2);
    vec sin_= sin(an2);
    mat A(p->image_nums,3,fill::ones);

    for (int j = 0; j < p->image_nums; j++) {
        A(j,1) = cos_(j);
        A(j,2) = sin_(j);
    }

    mat x= solve(A,I.t(),solve_opts::fast);//
    mat x_t=x.t();

    vec Iun_=x_t.col(0) ;
    vec rho_=sqrt(square(x_t.col(1))+square(x_t.col(2))) ;
    rho_=rho_/Iun_;
    rho_.elem(arma::find(rho_>1)).ones(); //elements limited
    vec phi_2= 0.5*atan2(x_t.col(2),x_t.col(1));//
    vec phi_= mod(phi_2,datum::pi);
    vec phi_spec(p->spec_rows.n_elem,fill::zeros);
  /*  We have a mask so need to reshape the estimated quantities to the
     masked pixels only*/
    for (int32_t j = 0; j < p->mask_cols.n_elem ; ++j) {
        if(p->spec(p->mask_rows(j),p->mask_cols(j))){
            phi_(j)=phi_(j)+0.5*datum::pi;
            phi_(j)=phi_(j)-std::floor(phi_(j)/datum::pi)*datum::pi;
        }
        p->phi(p->mask_rows(j),p->mask_cols(j)) = phi_(j);
        p->Iun(p->mask_rows(j),p->mask_cols(j)) = Iun_(j);
        p->rho(p->mask_rows(j),p->mask_cols(j)) = rho_(j);
    }

}

/*
 RHO_DIFFUSE Compute zenith angle from DOP for diffuse pixels
   Uses closed form inversion of DOP expression
   Inputs:
      rho - matrix or vector containing degree of polarisation values
      n   - scalar containing index of refraction
   Output:
      theta - zenith angle estimates (same size as rho)
 */
void NfP::Polarized::rho_diffuse( NfP::nfp_parameters* p) {
    float n_2= pow(p->n,2);
    float n_3= pow(p->n,3);
    float n_4= pow(p->n,4);
    mat rho_2=square(p->rho);

    mat Sub_style1=1-rho_2;
    Sub_style1=p->rho%sqrt(Sub_style1);
    mat zero(size(p->rho),fill::zeros);
    mat one(size(p->rho),fill::ones);

    mat temp=sqrt((2*p->rho + 2*n_2*p->rho - 2*n_2 + n_4 + rho_2 + 4*n_2*rho_2 - n_4*rho_2
            - 4*n_3*Sub_style1 + 1) /(n_4*rho_2 + 2*n_4*p->rho+ n_4 + 6*n_2*rho_2 + 4*n_2*p->rho
                    - 2*n_2 + rho_2 + 2*p->rho + 1) ) ;
// To avoid complex result in case of numerical issues:（The result may be plural）
    mat x=real(temp);
    mat theta_1 = arma::min(real(temp),one);
    p->theta = acos(theta_1);
}

/*
 RHO_S Compute zenith angle from degree of polarisation for specular pix
   There is no closed form inversion for rho_s so this functions uses a look up table and interpolates
   Inputs:
      rho - matrix or vector containing degree of polarisation values
      spec - the specular pix of images
      n   - scalar containing index of refraction
   Output:
      theta - zenith angle estimates (same size as rho)
 */
void NfP::Polarized::rho_spec( NfP::nfp_parameters* p)
{
    vec theta_s = linspace<vec>(0, 0.5*datum::pi, 158) ;
    vec sin_theta = sin(theta_s) ;
    vec cos_theta = cos(theta_s) ;
    vec sin_theta_2 = square(sin_theta) ;
    vec sin_theta_4 = pow(sin_theta,4) ;
    float n_2=pow(p->n,2) ;

    vec rho_s = (2*sin_theta_2 % cos_theta % sqrt(n_2-sin_theta_2))/(n_2-sin_theta_2-n_2*sin_theta_2 +2*sin_theta_4) ;

    uvec maxpos = find(rho_s==max(rho_s)) ;
    theta_s.set_size(max(maxpos)) ;
    rho_s.set_size(max(maxpos)) ;

    //1D data interpolation
    vec rho_(p->spec_rows.n_elem,fill::zeros) ;
    for (int32_t i = 0; i <p->spec_rows.n_elem; ++i) {
        rho_(i)=p->rho(p->spec_rows(i),p->spec_cols(i)) ;
    }

    vec theta_(p->spec_rows.n_elem,fill::zeros) ;
    interp1(rho_s,theta_s,rho_,theta_) ;
//    double mean=arma::mean(theta_);
//    cout<<"theta_spec mean ="<<mean<<endl;
    for (int32_t i = 0; i <p->spec_rows.n_elem; ++i) {
        p->theta(p->spec_rows(i),p->spec_cols(i)) = theta_(i);
    }

}
void NfP::Polarized::polar2normal(mat& zenith,mat& azimuth,cube* normal)
{
    normal->slice(0) = sin(zenith) % cos(azimuth);
    normal->slice(1) = sin(zenith) % sin(azimuth);
    normal->slice(2) = cos(zenith);

}